#!/usr/bin/env python3
"""
增强版无人机导航系统模拟
包含PID控制器、MPC控制和A*路径规划
"""

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import math
import time
from queue import PriorityQueue
import logging
from logging.handlers import RotatingFileHandler
import csv
import os
from datetime import datetime
# 在文件头部（靠近其他 import）添加：
from mpl_toolkits.mplot3d.art3d import Poly3DCollection

# 在文件开头添加以下代码（在所有import之后）
plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置支持中文的字体
plt.rcParams['axes.unicode_minus'] = False  # 正确显示负号
import copy


# 设置日志记录
def setup_logging():
    # 创建日志目录
    if not os.path.exists('logs'):
        os.makedirs('logs')

    # 创建日志文件名（带时间戳）
    log_filename = f"logs/drone_navigation_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log"

    # 设置日志格式
    log_format = '%(asctime)s - %(levelname)s - %(message)s'

    # 配置根日志记录器
    logging.basicConfig(
        level=logging.INFO,
        format=log_format,
        handlers=[
            RotatingFileHandler(log_filename, maxBytes=10 * 1024 * 1024, backupCount=5),
            logging.StreamHandler()
        ]
    )

    return log_filename


# 创建CSV数据记录器
class DataLogger:
    def __init__(self):
        # 创建数据目录
        if not os.path.exists('data'):
            os.makedirs('data')

        # 创建CSV文件名（带时间戳）
        self.filename = f"data/drone_data_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"

        # 初始化CSV文件
        with open(self.filename, 'w', newline='') as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow([
                'timestamp', 'pos_x', 'pos_y', 'pos_z',
                'vel_x', 'vel_y', 'vel_z',
                'target_x', 'target_y', 'target_z',
                'error_x', 'error_y', 'error_z', 'total_error',
                'control_force_x', 'control_force_y', 'control_force_z'
            ])

    def log_data(self, timestamp, position, velocity, target, control_force):
        error = target - position
        total_error = np.linalg.norm(error)

        with open(self.filename, 'a', newline='') as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow([
                timestamp,
                position[0], position[1], position[2],
                velocity[0], velocity[1], velocity[2],
                target[0], target[1], target[2],
                error[0], error[1], error[2], total_error,
                control_force[0], control_force[1], control_force[2]
            ])


class PIDController:
    """PID控制器实现"""

    def __init__(self, kp, ki, kd, max_output=float('inf')):
        # 确保增益参数是numpy数组
        self.kp = np.array(kp)  # 比例系数
        self.ki = np.array(ki)  # 积分系数
        self.kd = np.array(kd)  # 微分系数
        self.max_output = max_output  # 输出限制

        self.integral = np.zeros(3)
        self.previous_error = np.zeros(3)
        self.previous_time = time.time()

        # 记录PID各项输出
        self.p_term = np.zeros(3)
        self.i_term = np.zeros(3)
        self.d_term = np.zeros(3)

    def compute(self, setpoint, current_value):
        """计算控制输出"""
        current_time = time.time()
        dt = current_time - self.previous_time
        if dt <= 0:
            return np.zeros(3)

        error = setpoint - current_value

        # 比例项
        self.p_term = self.kp * error

        # 积分项 (带有抗饱和)
        self.integral += error * dt

        # 积分限幅 - 对每个维度分别处理
        max_integral = np.zeros(3)
        for i in range(3):
            if self.ki[i] != 0:
                max_integral[i] = 5.0 / self.ki[i]
            else:
                max_integral[i] = float('inf')

        self.integral = np.clip(self.integral, -max_integral, max_integral)
        self.i_term = self.ki * self.integral

        # 微分项 (带有滤波)
        derivative = (error - self.previous_error) / dt
        self.d_term = self.kd * derivative

        # 计算总输出
        output = self.p_term + self.i_term + self.d_term

        # 限制输出
        output_norm = np.linalg.norm(output)
        if output_norm > self.max_output:
            output = output / output_norm * self.max_output

        # 更新状态
        self.previous_error = error
        self.previous_time = current_time

        # 记录PID调试信息
        logging.debug(f"PID - P: {self.p_term}, I: {self.i_term}, D: {self.d_term}, Output: {output}")

        return output


class MPCController:
    """改进版模型预测控制器"""

    def __init__(self, horizon=5, dt=0.1):
        self.horizon = horizon  # 预测步长
        self.dt = dt  # 时间步长

        # 权重矩阵
        self.Q = np.diag([10, 10, 10])  # 状态误差权重
        self.R = np.diag([0.1, 0.1, 0.1])  # 控制输入权重
        self.Qf = np.diag([20, 20, 20])  # 终端状态误差权重

        # 系统矩阵 (简化的双积分器模型)
        self.A = np.eye(3)
        self.B = np.eye(3) * dt

        # 初始化
        self.last_control = np.zeros(3)
        self.last_error = np.zeros(3)

    def compute(self, current_state, reference_trajectory):
        """计算控制输入"""
        # 简化的MPC实现 - 使用线性二次调节器(LQR)方法
        # 在实际应用中，这里需要求解更复杂的优化问题

        # 提取当前位置和速度
        current_pos = current_state[:3]
        current_vel = current_state[3:6]

        # 计算到参考轨迹的误差
        error = reference_trajectory[0] - current_pos

        # 简化的控制律 (基于LQR)
        # 这里使用一个简化的反馈控制律，实际MPC应该考虑多步预测
        Kp = 3.0  # 位置增益
        Kd = 2.0  # 速度增益

        control = Kp * error + Kd * (self.last_error - error) / self.dt

        # 限制控制输入
        control_norm = np.linalg.norm(control)
        if control_norm > 5.0:
            control = control / control_norm * 5.0

        self.last_error = error
        self.last_control = control

        # 记录MPC调试信息
        logging.debug(f"MPC - Error: {error}, Control: {control}")

        return control


class AStarPlanner:
    """改进版A*路径规划算法实现"""

    def __init__(self, grid_size=1.0):
        self.grid_size = grid_size

    def heuristic(self, a, b):
        """计算启发式函数（欧几里得距离）"""
        return math.sqrt((a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2 + (a[2] - b[2]) ** 2)

    def plan(self, start, goal, obstacles):
        """执行A*路径规划"""
        # 将连续坐标转换为网格坐标
        start_grid = (int(round(start[0] / self.grid_size)),
                      int(round(start[1] / self.grid_size)),
                      int(round(start[2] / self.grid_size)))
        goal_grid = (int(round(goal[0] / self.grid_size)),
                     int(round(goal[1] / self.grid_size)),
                     int(round(goal[2] / self.grid_size)))

        # 创建障碍物网格
        obstacle_grid = set()
        for obs in obstacles:
            # 考虑障碍物大小 (假设障碍物占据一个网格)
            obs_grid = (int(round(obs[0] / self.grid_size)),
                        int(round(obs[1] / self.grid_size)),
                        int(round(obs[2] / self.grid_size)))
            obstacle_grid.add(obs_grid)

        # A*算法
        frontier = PriorityQueue()
        frontier.put((0, start_grid))
        came_from = {}
        cost_so_far = {}
        came_from[start_grid] = None
        cost_so_far[start_grid] = 0

        # 可能的移动方向（26个方向）
        directions = []
        for dx in [-1, 0, 1]:
            for dy in [-1, 0, 1]:
                for dz in [-1, 0, 1]:
                    if dx == 0 and dy == 0 and dz == 0:
                        continue
                    directions.append((dx, dy, dz))

        found_path = False
        while not frontier.empty():
            current = frontier.get()[1]

            if current == goal_grid:
                found_path = True
                break

            for dx, dy, dz in directions:
                next_node = (current[0] + dx, current[1] + dy, current[2] + dz)

                if next_node in obstacle_grid:
                    continue

                # 检查边界
                if abs(next_node[0]) > 50 or abs(next_node[1]) > 50 or next_node[2] < 0 or next_node[2] > 20:
                    continue

                new_cost = cost_so_far[current] + math.sqrt(dx * dx + dy * dy + dz * dz)
                if next_node not in cost_so_far or new_cost < cost_so_far[next_node]:
                    cost_so_far[next_node] = new_cost
                    priority = new_cost + self.heuristic(goal_grid, next_node)
                    frontier.put((priority, next_node))
                    came_from[next_node] = current

        if not found_path:
            logging.warning("A*路径规划失败，使用直线路径")
            return [start, goal]

        # 重建路径
        path = []
        current = goal_grid
        while current != start_grid:
            path.append(current)
            if current not in came_from:
                logging.error("A*路径重建失败")
                return [start, goal]
            current = came_from[current]
        path.append(start_grid)
        path.reverse()

        # 将网格坐标转换回连续坐标
        continuous_path = []
        for p in path:
            continuous_path.append(np.array([
                p[0] * self.grid_size,
                p[1] * self.grid_size,
                p[2] * self.grid_size
            ]))

        # 路径平滑
        smoothed_path = self.smooth_path(continuous_path, obstacles)

        logging.info(f"A*规划完成，共{len(smoothed_path)}个航点")
        return smoothed_path

    def smooth_path(self, path, obstacles):
        """使用简单的路径平滑算法"""
        if len(path) <= 2:
            return path

        smoothed = [path[0]]
        i = 0

        while i < len(path) - 1:
            j = len(path) - 1
            while j > i + 1:
                # 确保传递numpy数组
                if self.is_line_clear(np.array(path[i]), np.array(path[j]), obstacles):
                    smoothed.append(path[j])
                    i = j
                    break
                j -= 1
            else:
                smoothed.append(path[i + 1])
                i += 1

        return smoothed

    def is_line_clear(self, start, end, obstacles):
        """检查两点之间是否有障碍物"""
        # 将元组转换为numpy数组
        start = np.array(start)
        end = np.array(end)

        # 简化的直线检测
        steps = 20
        for step in range(steps + 1):
            t = step / steps
            point = start + t * (end - start)

            # 检查是否靠近任何障碍物
            for obs in obstacles:
                if np.linalg.norm(point - obs) < 1.5:  # 障碍物半径+安全距离
                    return False

        return True


class Drone:
    def __init__(self, x=0, y=0, z=0, mass=1.0):
        # 初始状态
        self.position = np.array([x, y, z], dtype=float)
        self.velocity = np.array([0, 0, 0], dtype=float)
        self.acceleration = np.array([0, 0, 0], dtype=float)

        # 物理参数
        self.mass = mass
        self.max_thrust = 30.0  # 增加最大推力 (N)
        self.gravity = np.array([0, 0, -9.81])  # 重力加速度

        # 控制器 - 调整PID参数，特别是z方向的增益
        self.pid_controller = PIDController(
            kp=np.array([2.0, 2.0, 6.0]),  # z方向增益更高
            ki=np.array([0.1, 0.1, 0.1]),  # z方向积分增益
            kd=np.array([1.0, 1.0, 3.0]),  # z方向微分增益更高
            max_output=25.0
        )
        self.mpc_controller = MPCController()
        self.controller_type = "PID"  # 默认为PID控制器

        # 历史轨迹记录
        self.trajectory = []

        # 控制力记录
        self.control_forces = []

        # 步数计数器
        self.step_count = 0

        # 重力补偿 - 应该恰好抵消重力
        self.gravity_compensation = np.array([0, 0, 9.81 * self.mass])  # 用于抵消重力

    def set_controller(self, controller_type):
        """设置控制器类型"""
        self.controller_type = controller_type
        logging.info(f"切换控制器到: {controller_type}")

    def update(self, target_position, dt=0.1):
        """更新无人机状态"""
        self.step_count += 1

        # 根据控制器类型计算控制力
        if self.controller_type == "PID":
            # PID控制 - 添加重力补偿
            control_force = self.pid_controller.compute(target_position, self.position)
            # 只在z方向添加重力补偿
            control_force[2] += self.gravity_compensation[2]
        else:
            # MPC控制
            ref_trajectory = [target_position] * self.mpc_controller.horizon
            current_state = np.concatenate([self.position, self.velocity])
            control_force = self.mpc_controller.compute(current_state, ref_trajectory)
            # 只在z方向添加重力补偿
            control_force[2] += self.gravity_compensation[2]

        # 记录控制力
        self.control_forces.append(control_force.copy())

        # 计算总力 (控制力 + 重力)
        total_force = control_force + self.mass * self.gravity

        # 限制推力
        thrust_norm = np.linalg.norm(total_force)
        if thrust_norm > self.max_thrust:
            total_force = total_force / thrust_norm * self.max_thrust

        # 计算加速度 (F = ma)
        self.acceleration = total_force / self.mass

        # 更新速度和位置
        self.velocity += self.acceleration * dt
        self.position += self.velocity * dt

        # 确保高度不低于0
        if self.position[2] < 0:
            self.position[2] = 0
            self.velocity[2] = 0

        # 记录轨迹
        self.trajectory.append(self.position.copy())

        # 每10步记录一次详细信息
        if self.step_count % 10 == 0:
            error = target_position - self.position
            logging.info(
                f"步数 {self.step_count}: 位置={self.position}, 目标={target_position}, "
                f"误差={error}, 控制力={control_force}"
            )

        # 检查是否到达目标
        distance = np.linalg.norm(target_position - self.position)
        return distance < 0.5  # 增加容差范围

    def get_position(self):
        return self.position.copy()

    def get_velocity(self):
        return self.velocity.copy()

    def get_trajectory(self):
        return np.array(self.trajectory)

    def get_control_forces(self):
        return np.array(self.control_forces)


class NavigationSystem:
    def __init__(self):
        # 设置日志
        self.log_filename = setup_logging()
        self.data_logger = DataLogger()

        logging.info("初始化导航系统...")

        # 创建障碍物
        self.obstacles = [
            np.array([5, 5, 2]),
            np.array([8, 3, 3]),
            np.array([12, 8, 4]),
            np.array([15, 12, 5]),
            np.array([7, 10, 3]),
            np.array([10, 15, 4]),
            np.array([13, 5, 3]),
            np.array([18, 8, 4])
        ]

        # 起点和终点
        self.start = np.array([0, 0, 1])
        self.goal = np.array([20, 15, 6])

        # 路径规划器
        self.planner = AStarPlanner(grid_size=1.0)

        # 规划路径
        self.waypoints = self.plan_path()

        # 当前目标航点索引
        self.current_waypoint_idx = 1  # 从第一个航点开始

        # 创建无人机实例
        self.drone = Drone(
            x=self.start[0],
            y=self.start[1],
            z=self.start[2]
        )

        # 是否完成所有航点
        self.mission_complete = False

        # 开始时间
        self.start_time = time.time()

        logging.info(f"导航系统初始化完成，共{len(self.waypoints)}个航点")

    def plan_path(self):
        """使用A*算法规划路径"""
        logging.info("正在规划路径...")
        path = self.planner.plan(self.start, self.goal, self.obstacles)

        if not path:
            logging.warning("无法找到路径! 使用直线路径")
            return [self.start, self.goal]

        logging.info(f"路径规划完成，共{len(path)}个航点")
        return path

    def update(self, dt=0.1):
        """更新导航系统状态"""
        if self.mission_complete:
            return True

        # 获取当前目标航点
        target = self.waypoints[self.current_waypoint_idx]

        # 更新无人机位置
        reached = self.drone.update(target, dt)

        # 记录数据
        current_time = time.time() - self.start_time
        self.data_logger.log_data(
            current_time,
            self.drone.get_position(),
            self.drone.get_velocity(),
            target,
            self.drone.control_forces[-1] if self.drone.control_forces else np.zeros(3)
        )

        # 如果到达当前航点，则转向下一个航点
        if reached:
            logging.info(f"到达航点 {self.current_waypoint_idx}: {target}")
            self.current_waypoint_idx += 1

            # 检查是否完成所有航点
            if self.current_waypoint_idx >= len(self.waypoints):
                logging.info("任务完成!")
                self.mission_complete = True
                return True

        return False

    def get_drone_position(self):
        return self.drone.get_position()

    def get_waypoints(self):
        return self.waypoints

    def get_trajectory(self):
        return self.drone.get_trajectory()

    def get_control_forces(self):
        return self.drone.get_control_forces()

    def get_obstacles(self):
        return self.obstacles

    def set_controller(self, controller_type):
        """设置控制器类型"""
        self.drone.set_controller(controller_type)

    def print_summary(self):
        """打印任务摘要"""
        trajectory = self.get_trajectory()
        total_distance = 0
        for i in range(1, len(trajectory)):
            total_distance += np.linalg.norm(trajectory[i] - trajectory[i - 1])

        logging.info(f"任务摘要: 总飞行距离={total_distance:.2f}m, 航点数={len(self.waypoints)}")
        logging.info(f"日志文件: {self.log_filename}")


class Simulation:
    def __init__(self, controller_type="PID"):
        # 创建导航系统
        self.nav_system = NavigationSystem()
        self.nav_system.set_controller(controller_type)

        # 创建图形
        self.fig = plt.figure(figsize=(15, 10))

        # 创建3D子图
        self.ax = self.fig.add_subplot(111, projection='3d')

        # 设置图表标题和标签
        self.ax.set_title(f'无人机导航模拟 - {controller_type}控制器')
        self.ax.set_xlabel('X (m)')
        self.ax.set_ylabel('Y (m)')
        self.ax.set_zlabel('Z (m)')

        # 设置坐标轴范围
        self.ax.set_xlim(-2, 25)
        self.ax.set_ylim(-2, 20)
        self.ax.set_zlim(0, 10)

        # 绘制起点和终点
        start = self.nav_system.start
        goal = self.nav_system.goal
        self.ax.scatter(start[0], start[1], start[2], c='green', marker='o', s=100, label='起点')
        self.ax.scatter(goal[0], goal[1], goal[2], c='red', marker='o', s=100, label='终点')

        # 绘制航点
        waypoints = self.nav_system.get_waypoints()
        wp_x = [wp[0] for wp in waypoints]
        wp_y = [wp[1] for wp in waypoints]
        wp_z = [wp[2] for wp in waypoints]
        self.ax.plot(wp_x, wp_y, wp_z, 'r--', alpha=0.5, label='规划路径')

        # 绘制障碍物
        obstacles = self.nav_system.get_obstacles()
        for obs in obstacles:
            self.draw_cube(obs, 1.0, 'gray')

        # 初始化无人机位置标记
        self.drone_marker, = self.ax.plot([], [], [], 'bo', markersize=8, label='无人机')

        # 初始化轨迹线
        self.trajectory_line, = self.ax.plot([], [], [], 'b-', alpha=0.7, linewidth=2, label='实际轨迹')

        # 初始化当前目标标记
        self.target_marker, = self.ax.plot([], [], [], 'rx', markersize=10, label='当前目标')

        # 添加图例
        self.ax.legend()

        # 设置视角
        self.ax.view_init(elev=20, azim=45)

        # 添加信息文本框
        self.info_text = self.ax.text2D(0.02, 0.95, "", transform=self.ax.transAxes)

        # 动画对象
        self.animation = None

        # 帧计数器
        self.frame_count = 0

    def draw_cube(self, center, size, color):
        """绘制立方体障碍物（使用 Poly3DCollection，避免 plot_trisurf 的三角剖分问题）"""
        # 计算立方体的顶点
        r = size / 2.0
        vertices = [
            [center[0] - r, center[1] - r, center[2] - r],
            [center[0] + r, center[1] - r, center[2] - r],
            [center[0] + r, center[1] + r, center[2] - r],
            [center[0] - r, center[1] + r, center[2] - r],
            [center[0] - r, center[1] - r, center[2] + r],
            [center[0] + r, center[1] - r, center[2] + r],
            [center[0] + r, center[1] + r, center[2] + r],
            [center[0] - r, center[1] + r, center[2] + r]
        ]

        # 定义立方体的面 (每个面为4个顶点)
        faces = [
            [0, 1, 2, 3],
            [4, 5, 6, 7],
            [0, 1, 5, 4],
            [2, 3, 7, 6],
            [0, 3, 7, 4],
            [1, 2, 6, 5]
        ]

        # 使用 Poly3DCollection 绘制每个面
        poly_verts = []
        for face in faces:
            poly_verts.append([vertices[i] for i in face])

        collection = Poly3DCollection(poly_verts, facecolors=color, linewidths=0.5, edgecolors='k', alpha=0.25)
        self.ax.add_collection3d(collection)

    def update_plot(self, frame):
        """更新动画帧"""
        self.frame_count += 1

        # 更新导航系统
        complete = self.nav_system.update(dt=0.1)

        # 获取无人机当前位置
        drone_pos = self.nav_system.get_drone_position()

        # 获取当前目标
        if self.nav_system.current_waypoint_idx < len(self.nav_system.waypoints):
            target_pos = self.nav_system.waypoints[self.nav_system.current_waypoint_idx]
        else:
            target_pos = self.nav_system.waypoints[-1]

        # 更新无人机位置标记
        self.drone_marker.set_data([drone_pos[0]], [drone_pos[1]])
        self.drone_marker.set_3d_properties([drone_pos[2]])

        # 更新目标位置标记
        self.target_marker.set_data([target_pos[0]], [target_pos[1]])
        self.target_marker.set_3d_properties([target_pos[2]])

        # 更新轨迹线
        trajectory = self.nav_system.get_trajectory()
        if len(trajectory) > 1:
            self.trajectory_line.set_data(trajectory[:, 0], trajectory[:, 1])
            self.trajectory_line.set_3d_properties(trajectory[:, 2])

        # 更新信息文本
        info = f"帧: {self.frame_count}\n位置: {drone_pos}\n目标: {target_pos}\n误差: {np.linalg.norm(target_pos - drone_pos):.2f}m"
        if complete:
            info += "\n状态: 任务完成!"
        self.info_text.set_text(info)

        # 如果任务完成，停止动画
        if complete:
            self.animation.event_source.stop()
            self.nav_system.print_summary()
            logging.info("模拟完成")

        return self.drone_marker, self.trajectory_line, self.target_marker, self.info_text

    def run(self):
        """运行模拟"""
        logging.info("开始无人机导航模拟...")
        logging.info(f"控制器类型: {self.nav_system.drone.controller_type}")
        logging.info(f"起点: {self.nav_system.start}")
        logging.info(f"终点: {self.nav_system.goal}")
        logging.info(f"障碍物数量: {len(self.nav_system.get_obstacles())}")

        # 创建动画
        self.animation = FuncAnimation(
            self.fig, self.update_plot, frames=1000,
            interval=50, blit=False, repeat=False
        )

        plt.show()


# 比较不同控制器的性能
def compare_controllers():
    """比较PID和MPC控制器的性能"""
    logging.info("比较不同控制器的性能...")

    # 运行PID控制器模拟
    logging.info("\n=== PID控制器 ===")
    sim_pid = Simulation(controller_type="PID")
    sim_pid.run()

    # 运行MPC控制器模拟
    logging.info("\n=== MPC控制器 ===")
    sim_mpc = Simulation(controller_type="MPC")
    sim_mpc.run()


# 运行模拟
if __name__ == "__main__":
    # 运行单个模拟
    sim = Simulation(controller_type="PID")  # 尝试改为"MPC"
    sim.run()

    # 或者比较不同控制器
    # compare_controllers()